package com.lw.leetcode.arr.b;

import java.util.Arrays;

/**
 * Created with IntelliJ IDEA.
 * 6017. 向数组中追加 K 个整数
 *
 * @author liw
 * @version 1.0
 * @date 2022/3/6 20:01
 */
public class MinimalKSum {


    public static void main(String[] args) {
        MinimalKSum test = new MinimalKSum();

        // 5
//        int[] arr = {1, 4, 5};
//        int k = 2;

        // 25
//        int[] arr = {5,6};
//        int k = 6;

        // 794
//        int[] arr = {96,44,99,25,61,84,88,18,19,33,60,86,52,19,32,47,35,50,94,17,29,98,22,21,72,100,40,84};
//        int k = 35;


        // 250
        int[] arr = {93, 44, 49, 45, 93, 52, 6, 7, 88, 70, 86, 15, 38, 86, 86, 95, 8, 62, 13, 84, 26, 16, 33, 85, 7, 62, 55, 50, 77, 10, 76, 10, 35, 67, 19, 12, 24, 39, 76, 37};
        int k = 17;

        long sum = test.minimalKSum(arr, k);
        System.out.println(sum);
    }

    public long minimalKSum(int[] nums, int k) {
        int length = nums.length;
        Arrays.sort(nums);
        long sum = 0L;
        int n = 0;
        for (int i = 0; i < length; i++) {
            if (k == 0) {
                return sum;
            }
            long c = nums[i] - n - 1L;
            if (c <= 0) {
                n = nums[i];
                continue;
            }
            if (c <= k) {
                k -= c;
                sum += (2L * n + 1 + c) * c / 2;
            } else {
                sum += (2L * n + 1 + k) * k / 2;
                k = 0;
            }
            n = nums[i];
        }
        if (k > 0) {
            sum += (2L * nums[length - 1] + 1 + k) * k / 2;
        }
        return sum;
    }

}
